#include "log.h"
#include "net.h"
#include "dial.h"

struct net_packet {
    struct iphdr iph;
    struct icmp icp;
};

void tr_result_print(tr_result_t* result)
{
    int i = 0;
    for (i = 0; i < result->count; i++) {
        log_trace("ip:%s, latency:%.2f, ttl:%d",
            inet_ntoa(*((struct in_addr*)&result->ret[i].ip)), result->ret[i].latency,
            result->ret[i].ttl);
    }
    return;
}

void send_probe_pkt(int fd, char* content, int len, struct sockaddr_in* addr, int count)
{
    int i = 0;

    for (i = 0; i < count; i++) {
        sendto(fd, content, len, 0, (struct sockaddr*)addr, sizeof(*addr));
    }
}

int read_response_pkt(int fd, char* buf, int len, int opt, uint16_t id, uint16_t seq, uint64_t snd_time, tr_ret_t* tr_ret)
{
    struct sockaddr_in peer_addr;
    int addr_len = sizeof(peer_addr);
    struct net_packet* pkt;
    struct icmp* icp;
    struct iphdr *iph;
    struct udphdr* udp_hdr;
    uint32_t tm;
    struct timeval now;
    time_t r_now, start_tm;
    

    start_tm = time(NULL);
    while (1) {
        r_now = time(NULL);
        if (r_now - start_tm > 3) {
            return 0;
        }
        int ret = recvfrom(fd, buf, len, 0, (struct sockaddr*)&peer_addr, &addr_len);
        if (ret < 0) {
            if (errno == EAGAIN)
                return 0;
            log_syserr("traceroute rcv error.");
            return 0;
        }
        gettimeofday(&now, NULL);
        iph = (struct iphdr *)buf;
        icp = (struct icmp *)((char *)iph + iph->ihl*4);

        switch (icp->icmp_type) {
        case ICMP_TIME_EXCEEDED:
            if (opt == OPT_PRO_ICMP) {
                iph = (struct iphdr *)((char *)icp + ICMP_MINLEN);
                icp = (struct icmp *)((char *)iph + iph->ihl *4);
                if (!(icp->icmp_id == id && icp->icmp_seq == seq)) {
                    continue;
                }
            } else if (opt == OPT_PRO_UDP) {
                iph = (struct iphdr *)((char *)icp + ICMP_MINLEN);
                udp_hdr = (struct udphdr *)((char *)iph + iph->ihl *4);
                // log_debug("udp_hdr->source :%d, id:%d", udp_hdr->source ,id);
                if (udp_hdr->source != id) {
                    continue;
                }
            }
            tm = (now.tv_sec * 1000000ULL + now.tv_usec - snd_time);
            tr_ret->latency = tm / 1000.0;
            tr_ret->ip = peer_addr.sin_addr.s_addr;
            log_trace("peer_addr:%s, latency:%.2f", inet_ntoa(peer_addr.sin_addr), tr_ret->latency);
            return 0;
            break;
        case ICMP_ECHOREPLY:
            if (icp->icmp_id != id) {
                continue;
            }
            tm = (now.tv_sec * 1000000ULL + now.tv_usec - snd_time);
            tr_ret->latency = tm / 1000.0;
            tr_ret->ip = peer_addr.sin_addr.s_addr;
            log_trace("peer_addr:%s, latency:%.2f", inet_ntoa(peer_addr.sin_addr), tr_ret->latency);
            return 1;
        case ICMP_DEST_UNREACH:
            if (icp->icmp_code == ICMP_PORT_UNREACH) {
                iph = (struct iphdr *)((char *)icp + ICMP_MINLEN);
                udp_hdr = (struct udphdr *)((char *)iph + iph->ihl *4);
                if (udp_hdr->source != id) {
                    continue;
                }
                tm = (now.tv_sec * 1000000ULL + now.tv_usec - snd_time);
                tr_ret->latency = tm / 1000.0;
                tr_ret->ip = peer_addr.sin_addr.s_addr;
                log_trace("peer_addr:%s, latency:%.2f", inet_ntoa(peer_addr.sin_addr), tr_ret->latency);
                return 1;
            }
            break;
        default:
            break;
        }
    }

    return 1;
}

void* do_traceroute(void* arg)
{
    dial_arg_t* dial_arg = (dial_arg_t*)arg;
    dial_result_t* dial_result = calloc(1, sizeof(dial_result_t));
    tr_arg_t* tr_arg = &dial_arg->in.tr_arg;
    tr_result_t* tr_result = &dial_result->out.tr_result;

    int raw_fd = 0, snd_fd = 0, hop = 0;
    char buf[1500] = { 0 }, rbuf[1500] = { 0 }, *content = NULL;
    struct icmp* icp;
    struct udphdr* udp_hdr;
    struct sockaddr_in addr, peer_addr;
    int icmpseq = 1, ttl = 1, blksize = 0;
    uint16_t id = 0, src_port;
    uint64_t snd_time;
    struct timeval tm;

    blksize = tr_arg->blocksize ? tr_arg->blocksize : 64;
    raw_fd = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
    addr.sin_family = AF_INET;
    addr.sin_port = tr_arg->port ? tr_arg->port : TRACEROUTE_UDP_DEF_PORT;
    addr.sin_addr.s_addr = tr_arg->ip;
    if (tr_arg->opt == OPT_PRO_ICMP) {
        snd_fd = raw_fd;
        icp = (struct icmp*)buf;
        content = (char*)icp;
        fill_icmp_packet(icp, ICMP_ECHO, icmpseq);
    } else {
        snd_fd = socket(AF_INET, SOCK_RAW, IPPROTO_UDP);
        udp_hdr = (struct udphdr*)buf;
        content = (char*)udp_hdr;
        src_port = port_random(0);
    }

    sock_rcvtmo(raw_fd, 2, 0);
    hop = tr_arg->max_hop ? tr_arg->max_hop : TR_MAX_HOP;
    while (hop--) {
        if (tr_arg->opt == OPT_PRO_ICMP) {
            icmpseq++;
            icmp_set_seq(icp, icmpseq);
            id = icp->icmp_id;
            snd_time = *((uint32_t*)&icp->icmp_dun);
        } else {
            id = ++src_port;
            fill_udp_packet(udp_hdr, src_port, tr_arg->port, blksize);
            snd_time = monotonic_us();
        }
        set_ttl(snd_fd, ttl);
        send_probe_pkt(snd_fd, content, blksize, &addr, 5);

        tr_result->count++;
        tr_result->ret[ttl - 1].ttl = ttl;
        if (read_response_pkt(raw_fd, rbuf, sizeof(rbuf), tr_arg->opt, id, icmpseq, snd_time, &tr_result->ret[ttl - 1]))
            break;
        ttl++;
    }
    close(snd_fd);
    close(raw_fd);
    return (void*)dial_result;
}